﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Drawing;

namespace LocalizationTest
{
  public class BeliefProbabilityDistribution
  {
    Dictionary<WorldState, double> statesWithProbability;

    public Dictionary<WorldState, double> StatesWithProbability
    {
      get { return statesWithProbability; }
    }

    public BeliefProbabilityDistribution()
    {
      statesWithProbability = new Dictionary<WorldState,double>();
    }

    public void Draw(string drawingID)
    {
      foreach (WorldState worldState in statesWithProbability.Keys)
      {
        worldState.Draw(statesWithProbability[worldState], drawingID);
      }
    }

    public void AddValue(WorldState worldState, double probability)
    {
      if (statesWithProbability.ContainsKey(worldState))
      {
        statesWithProbability[worldState] += probability;
      }
      else
      {
        statesWithProbability.Add(worldState, probability);
      }
    }

    public void AddValueUnsave(bool init, WorldState worldState, double probability)
    {
      if (init)
      {
        statesWithProbability.Add(worldState, probability);
      }
      else
      {
        statesWithProbability[worldState] += probability;
      }
    }

    public void InitWithEqualDistribution()
    {
      foreach (WorldState key in new Dictionary<WorldState, double>(statesWithProbability).Keys)
      {
        statesWithProbability[key] = 1.0 / statesWithProbability.Values.Count;
      }
    }

    public void Normalize()
    {
      double sum = 0;
      foreach(double value in statesWithProbability.Values)
      {
        sum += value;
      }
      foreach(WorldState key in new Dictionary<WorldState, double>(statesWithProbability).Keys)
      {
        if (sum == 0)
        {
          statesWithProbability[key] = 1.0 / statesWithProbability.Values.Count;
        }
        else statesWithProbability[key] /= sum;
      }
    }

    public WorldState GetMaximum()
    {
      WorldState worldStateWithHighestProbability = null;
      double maximum = 0;

      foreach (WorldState key in new Dictionary<WorldState, double>(statesWithProbability).Keys)
      {
        if (statesWithProbability[key] >= maximum)
        {
          maximum = statesWithProbability[key];
          worldStateWithHighestProbability = key;
        }
      }
      return worldStateWithHighestProbability;
    }

    public void UpdateBeliefWithAction(MotionModel motionModel, ActionDescription action)
    {
      BeliefProbabilityDistribution newBelief = new BeliefProbabilityDistribution();
      int counter = 0;
      foreach (WorldState oldWorldState in this.StatesWithProbability.Keys)
      {
        foreach (WorldState newWorldState in this.StatesWithProbability.Keys)
        {
          if (counter == 0 || motionModel.StatesAreNearBy(oldWorldState, newWorldState))
          {
            newBelief.AddValueUnsave(
              counter == 0,
              newWorldState,
              this.StatesWithProbability[oldWorldState] * motionModel.GetStateTransitionProbability(oldWorldState, newWorldState, action)
              );
          }
        }
        counter++;
      }
      statesWithProbability = newBelief.statesWithProbability;
    }
  }
}
